import tensorflow as tf
from .network import BottleNeck, h_swish, ConvBlock


class MobileNetV3Small(tf.keras.Model):
    def __init__(self, num_classes=10, activation=tf.nn.softmax):
        super(MobileNetV3Small, self).__init__()

        self.head = ConvBlock(filters=16, kernel_size=(3, 3), strides=2, padding='same', use_bias=False, activation=h_swish)
        self.bneck1 = BottleNeck(in_size=16, exp_size=16, out_size=16, s=2, is_se_existing=True, NL="RE", k=3)
        self.bneck2 = BottleNeck(in_size=16, exp_size=72, out_size=24, s=2, is_se_existing=False, NL="RE", k=3)
        self.bneck3 = BottleNeck(in_size=24, exp_size=88, out_size=24, s=1, is_se_existing=False, NL="RE", k=3, name='feature_1_8')
        self.bneck4 = BottleNeck(in_size=24, exp_size=96, out_size=40, s=2, is_se_existing=True, NL="HS", k=5)
        self.bneck5 = BottleNeck(in_size=40, exp_size=240, out_size=40, s=1, is_se_existing=True, NL="HS", k=5)
        self.bneck6 = BottleNeck(in_size=40, exp_size=240, out_size=40, s=1, is_se_existing=True, NL="HS", k=5)
        self.bneck7 = BottleNeck(in_size=40, exp_size=120, out_size=48, s=1, is_se_existing=True, NL="HS", k=5)
        self.bneck8 = BottleNeck(in_size=48, exp_size=144, out_size=48, s=1, is_se_existing=True, NL="HS", k=5, name="feature_1_16")
        self.bneck9 = BottleNeck(in_size=48, exp_size=288, out_size=96, s=2, is_se_existing=True, NL="HS", k=5)
        self.bneck10 = BottleNeck(in_size=96, exp_size=576, out_size=96, s=1, is_se_existing=True, NL="HS", k=5)
        self.bneck11 = BottleNeck(in_size=96, exp_size=576, out_size=96, s=1, is_se_existing=True, NL="HS", k=5)

        self.last_conv1 = ConvBlock(filters=576, kernel_size=(1, 1), use_bias=False, activation=h_swish)
        self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(7, 7),
                                                        strides=1)
        self.last_conv2 = ConvBlock(filters=1280, kernel_size=(1, 1), use_bias=False, activation=h_swish, use_bn=False)
        self.last_conv3 = ConvBlock(filters=num_classes, kernel_size=(1, 1), use_bias=False, activation=activation, use_bn=False)

    def call(self, inputs, training=None, mask=None):
        x = self.head(inputs, training=training)

        x = self.bneck1(x, training=training)
        x = self.bneck2(x, training=training)
        x = self.bneck3(x, training=training)
        x = self.bneck4(x, training=training)
        x = self.bneck5(x, training=training)
        x = self.bneck6(x, training=training)
        x = self.bneck7(x, training=training)
        x = self.bneck8(x, training=training)
        x = self.bneck9(x, training=training)
        x = self.bneck10(x, training=training)
        x = self.bneck11(x, training=training)

        x = self.last_conv1(x, training=training)
        x = self.avgpool(x)
        x = self.last_conv2(x)
        x = self.last_conv3(x)
        x = tf.squeeze(x, axis=[1, 2])
        return x

if __name__ == '__main__':
    model = MobileNetV3Small()
    model.build(input_shape=(None, 224, 224, 3))
    model.summary()
    x = tf.random.normal((2, 224, 224, 3))
    y = model(x)
    print(y.shape)
